# -*- coding: utf-8 -*-
#############################################################################################
#############################################################################################
#
#   Open Source License/Disclaimer, Forecast Systems Laboratory NOAA/OAR/GSD,
#   325 Broadway Boulder, CO 80305
#
#   This software is distributed under the Open Source Definition, which may be
#   found at http://www.opensource.org/.
#   In particular, redistribution and use in source and binary forms, with or
#   without modification, are permitted provided that the following conditions are met:
#
#   - Redistributions of source code must retain this notice, this list of
#     conditions and the following disclaimer.
#
#   - Redistributions in binary form must provide access to this notice, this
#     list of conditions and the following disclaimer, and the underlying source code.
#
#   - All modifications to this software must be clearly documented, and are
#     solely the responsibility of the agent making the modifications.
#
#   - If significant modifications or enhancements are made to this software,
#     the GSD Software Policy Manager (softwaremgr.fsl@noaa.gov) should be notified.
#
#   THIS SOFTWARE AND ITS DOCUMENTATION ARE IN THE PUBLIC DOMAIN AND ARE
#   FURNISHED "AS IS." THE AUTHORS, THE UNITED STATES GOVERNMENT, ITS INSTRUMENTALITIES,
#   OFFICERS, EMPLOYEES, AND AGENTS MAKE NO WARRANTY, EXPRESS OR IMPLIED, AS TO
#   THE USEFULNESS OF THE SOFTWARE AND DOCUMENTATION FOR ANY PURPOSE. THEY ASSUME
#   NO RESPONSIBILITY (1) FOR THE USE OF THE SOFTWARE AND DOCUMENTATION; OR (2) TO PROVIDE
#   TECHNICAL SUPPORT TO USERS.
#############################################################################################
#############################################################################################

import matplotlib.pyplot as plt
from matplotlib import cm
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.ticker import LinearLocator, FormatStrFormatter


#3D Plot
def plot_3d(X, Y, Z, name):
    fig = plt.figure()
    ax = fig.gca(projection='3d')               # to work in 3d

    surf = ax.plot_surface(X[1:-1,1:-1]/1000., Y[1:-1,1:-1]/1000., Z[1:-1,1:-1], rstride=1, cstride=1, \
            cmap=plt.get_cmap('jet'), vmax=10., vmin=-3., linewidth=0, antialiased=False)
    plt.axis([X[1:-1,1:-1].min()/1000., X[1:-1,1:-1].max()/1000., Y[1:-1,1:-1].min()/1000., Y[1:-1,1:-1].max()/1000.])
    ax.set_zlim(-3., 10.)

    ax.zaxis.set_major_locator(LinearLocator(10))
    ax.zaxis.set_major_formatter(FormatStrFormatter('%.02f'))

    fig.colorbar(surf, shrink=0.5, aspect=5)
    plt.title(name)
    plt.show()

    
def plot_2d(X, Y, Z, name):
    plt.figure()
    plt.imshow(Z[1:-1,1:-1], cmap=plt.get_cmap('jet'),vmax=10.,vmin=-3.)
    plt.colorbar()
    # set the limits of the plot to the limits of the data
    #plt.axis([X[1:-1,1:-1].min()/1000., X[1:-1,1:-1].max()/1000., Y[1:-1,1:-1].min()/1000., Y[1:-1,1:-1].max()/1000.])
    plt.title(name)
    plt.show()

